# -*- coding: utf-8 -*-

import logging
import sys
import threading
from datetime import datetime
from time import time

import pika
from pika import credentials

DURABLE_TIME = 50


class MQHandler(logging.Handler):
    """
    把日志打印到RabbitMQ服务端
    """

    def __init__(self, host, port=None, virtual_host=None, username=None, password=None, exchange='log'):
        logging.Handler.__init__(self)
        self.connection_para = dict(host=host)
        if port:
            self.connection_para['port'] = port
        if virtual_host:
            self.connection_para['virtual_host'] = virtual_host
        if username and password:
            self.connection_para['credentials'] = credentials.PlainCredentials(username, password)
        self.exchange = exchange
        self.connection, self.channel = None, None

        # 只声明一次exchange
        self.is_exchange_declared = False
        # 因为pika的线程安全问题, 上锁序列化发送过程
        self.emit_lock = threading.RLock()
        # 初始化连接
        self.connect()

    def connect(self):
        """
        连接rabbitMq, 使用topic类型的exchange
        """
        # 把pika的日志打到stdout, 以免在connect时发生RecursionError
        pika_logger = logging.getLogger('pika')
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter('%(asctime)s - %(module)s - %(levelname)s - %(message)s'))
        pika_logger.addHandler(handler)
        pika_logger.propagate = False
        pika_logger.setLevel(logging.WARNING)
        # 禁掉heartbeat.
        self.connection_para['heartbeat_interval'] = 0
        self.connection = pika.BlockingConnection(pika.ConnectionParameters(**self.connection_para))
        self.channel = self.connection.channel()
        sys.stdout.write('%s - stdout - [mq] Connect success.\n' % datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
        if not self.is_exchange_declared:
            self.channel.exchange_declare(exchange=self.exchange, exchange_type='topic', durable=True, auto_delete=False)
            sys.stdout.write(
                '%s - stdout - [mq] Declare exchange success.\n' % datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
            self.is_exchange_declared = True

    def emit(self, record):
        start_time = time()
        self.emit_lock.acquire()
        try:
            if not self.connection or not self.channel:
                self.connect()
            routing_key = "{name}.{level}".format(name=record.name, level=record.levelname)
            self.channel.basic_publish(exchange=self.exchange,
                                       routing_key=routing_key,
                                       body=self.format(record),
                                       properties=pika.BasicProperties(
                                           delivery_mode=2
                                       ))

        except Exception:
            # 如果connection关闭, 赋值为None,方便下次重连
            self.channel, self.connection = None, None
            self.handleError(record)
        finally:
            self.emit_lock.release()
            elapsed_time = (time() - start_time) * 1000
            if elapsed_time > DURABLE_TIME:
                sys.stdout.write('%s - stdout - [mq] Slow publishing, use time: %sms, durable time: %sms.\n'
                                 % (datetime.now().strftime("%Y-%m-%d %H:%M:%S"), elapsed_time, DURABLE_TIME))

    def close(self):
        """
        退出清理
        """
        self.acquire()
        try:
            logging.Handler.close(self)
            if self.channel:
                self.channel.close()
            if self.connection:
                self.connection.close()
            sys.stdout.write('%s - stdout - [mq] Clean up success.\n' % datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
        finally:
            self.release()
